import numpy as np

def loadDataSet(fileName):
    """
    1.Python strip() 方法用于移除字符串头尾指定的字符（默认为空格或换行符）或字符序列。
     注意：该方法只能删除开头或是结尾的字符，不能删除中间部分的字符。
    2.map() 会根据提供的函数对指定序列做映射。
     第一个参数 function 以参数序列中的每一个元素调用 function 函数，返回包含每次 function 函数返回值的新列表。

    函数说明:加载数据
    Parameters:
        fileName - 文件名
    Returns:
        dataMat - 数据矩阵
    """
    dataMat = []
    fr = open(fileName)
    for line in fr.readlines():
        curLine = line.strip().split("\t")
        fltLine = list(map(float, curLine))
        dataMat.append(fltLine)
    return dataMat


def binSplitDataSet(dataSet, feature, value):
    """
        函数说明:根据特征切分数据集合
        Parameters:
            dataSet - 数据集合
            feature - 带切分的特征
            value - 该特征的值
        Returns:
            mat0 - 切分的数据集合0
            mat1 - 切分的数据集合1
        """
    mat0 = dataSet[np.nonzero(dataSet[:, feature] > value)[0], :]
    mat1 = dataSet[np.nonzero(dataSet[:, feature] <= value)[0], :]
    return mat0, mat1


def regLeaf(dataSet):
    """
        函数说明:生成叶结点
        Parameters:
            dataSet - 数据集合
        Returns:
            目标变量的均值
    """
    return np.mean(dataSet[:, -1]) # 数据集列表最后一列特征值的均值作为叶节点返回

#    -----------------------模型树部分-------------------------
def regErr(dataSet):
    """
    函数说明:误差估计函数
    Parameters:
        dataSet - 数据集合
    Returns:
        目标变量的总方差
    """
    return np.var(dataSet[:,-1]) * np.shape(dataSet)[0]


def linearSolve(dataSet):
    m, n = np.shape(dataSet)
    X = np.mat(np.ones((m, n))) # 自变量
    Y = np.mat(np.ones((m, 1))) # 因变量
    X[:, 1:n] = dataSet[:, 0:n-1] # 样本数据集合
    Y = dataSet[:,-1] # 标签
    xTx = X.T * X # 线性模型 求解
    if np.linalg.det(xTx) == 0.0: # 判断矩阵是否可逆
        raise NameError('行列式值为零,不能计算逆矩阵，可适当增加ops的第二个值')
    ws = xTx.I * (X.T * Y)
    return ws, X, Y


def modelLeaf(dataSet): # 模型叶子节点
    ws, X, Y = linearSolve(dataSet)
    return ws


def modelErr(dataSet): # 计算模型误差
    ws, X, Y = linearSolve(dataSet)
    yHat = X * ws # 预测值
    return np.sum(np.power(Y - yHat, 2))


def chooseBestSplit(dataSet, leafType=regLeaf, errType=regErr, ops=(1, 4)):
    """
        函数说明:找到数据的最佳二元切分方式函数
        Parameters:
            dataSet - 数据集合
            leafType - 生成叶结点
            regErr - 误差估计函数
            ops - 用户定义的参数构成的元组
        Returns:
            bestIndex - 最佳切分特征
            bestValue - 最佳特征值
        """
    tolS = ops[0]; tolN = ops[1] # tolS允许的误差下降值,tolN切分的最少样本数
    # 如果当前所有值(剩余的特征值)相等,则退出。(根据set的特性)
    if len(set(dataSet[:, -1].T.tolist()[0])) == 1: # tolist:将数组或者矩阵转换成列表
        return None, leafType(dataSet)
    m, n = np.shape(dataSet)
    # 默认最后一个特征为最佳切分特征,计算其误差估计,计算未切分前数据集的误差,之后会遍历每个特征寻找最佳切分
    S = errType(dataSet)
    # 初始化最小误差；最佳切分特征索引(特征)；最佳切分特征值
    bestS = float("inf"); bestIndex = 0; bestValue = 0
    # 遍历所有特征，n-1是因为最后一列为默认的最佳切分，前面已经计算过了
    for featIndex in range(n-1):
        # 遍历所有特征值  matrix矩阵名.A代表将矩阵转化为array数组类型 因为dataSet是二维数组所以转置的时候会有2个中括号，所以加个[0],相当于去电一个[]
        for splitVal in set(dataSet[:, featIndex].T.A.tolist()[0]):
            mat0, mat1 = binSplitDataSet(dataSet, featIndex, splitVal) #根据特征和特征值切分数据集
            if (np.shape(mat0)[0] < tolN) or (np.shape(mat1)[0] < tolN): #如果数据少于tolN,则退出
                continue
            # 计算切分后的误差，即均方差和
            newS = errType(mat0) + errType(mat1)
            # 如果误差估计更小,则更新特征索引值和特征值
            if newS < bestS:
                bestIndex = featIndex
                bestValue = splitVal
                bestS = newS
    # 如果切分后比切分前误差下降值未达到tolS,不需切分，直接返回目标变量均值作为叶节点
    if (S - bestS) < tolS:
        return None, leafType(dataSet)
    # 根据最佳的切分特征和特征值切分数据集合,检查最佳特征及特征值是否满足不切分条件
    mat0, mat1 = binSplitDataSet(dataSet, bestIndex, bestValue)
    if (np.shape(mat0)[0] < tolN) or (np.shape(mat1)[0] < tolN):
        return None, leafType(dataSet)
    # 返回最佳切分特征和特征值
    return bestIndex, bestValue


def createTree(dataSet, leafType=regLeaf, errType=regErr, ops=(1, 4)):
    """
       函数说明:树构建函数
       Parameters:
           dataSet - 数据集合
           leafType - 建立叶结点的函数
           errType - 误差计算函数
           ops - 包含树构建所有其他参数的元组
       Returns:
           retTree - 构建的回归树

       """
    # 选择最佳切分特征和特征值
    feat, val = chooseBestSplit(dataSet, leafType, errType, ops)
    # 如果特征为none，直接返回叶节点值(特征值)
    if feat == None:
        return val
    # 回归树：字典类型
    retTree = {}
    # 树字典的一个元素是切分的最佳特征 第二个元素是最佳特征对应的最佳切分特征值
    retTree["spInd"] = feat
    retTree["spVal"] = val
    # 分成左数据集和右数据集
    lSet, rSet = binSplitDataSet(dataSet, feat, val)
    # 第三个元素是树的左分支，通过lSet子集递归生成左子树
    # 第四个元素是树的右分支，通过rSet子集递归生成右子树
    retTree["left"] = createTree(lSet, leafType, errType, ops)
    retTree["right"] = createTree(rSet, leafType, errType, ops)
    return retTree


#      -----------------------------后剪枝------------------------------
def isTree(obj):
    """
    用于测量输入变量是否是一棵树，是字典就是子树，不识字典就是叶结点即不是子树
    根据目标数据的存储类型是否为字典型，是返回true，否则返回false
    """
    return (type(obj).__name__ == 'dict')


def getMean(tree):
    """
    getMean()是个递归函数，从上往下遍历树直到叶节点为止
    :param tree:
    :return:
    """
    if isTree(tree['right']): # 树字典的右分支为字典类型（子树）：递归获得右子树的均值
        tree['right'] = getMean(tree['right'])
    if isTree(tree['left']):
        tree['left'] = getMean(tree['left'])
    return (tree['right'] + tree['left']) / 2.0 #递归直至找到两个叶节点，求二者的均值返回


def prune(tree, testData):
    """
    剪枝函数
    :param tree:待剪枝的树字典
    :param testData:用于剪枝的测试集
    :return:
    """
    if np.shape(testData)[0] == 0: # 测试集为空，直接对树相邻叶子结点进行求均值操作
        return getMean(tree)
    if isTree(tree['right']) or isTree(tree['left']): # 左右分支中有非叶子结点类型
        # 利用当前树的最佳切分点和特征值对测试集进行树构建过程
        lSet, rSet = binSplitDataSet(testData, tree["spInd"], tree["spVal"])
    # 左（右）分支非叶子结点，递归利用测试数据的左（右）子集对做分支剪枝
    if isTree(tree['left']):
        tree['left'] = prune(tree['left'], lSet)
    if isTree(tree['right']):
        tree['right'] = prune(tree['right'], rSet)
    # 如果左右分支都是叶子结点
    if not isTree(tree['left']) and not isTree(tree['right']):
        lSet, rSet = binSplitDataSet(testData, tree["spInd"], tree["spVal"])
        # 如果这两个叶节点不合并，计算误差，即（实际值-预测值）的平方和
        errorNoMerge = np.sum(np.power(lSet[:,-1] - tree['left'], 2)) + \
                       np.sum(np.power(rSet[:,-1] - tree['right'], 2))
        treeMean = (tree['left'] + tree['right']) / 2.0
        # 如果两个叶节点合并，计算合并后误差,即(真实值-合并后值）平方和
        errorMerge = np.sum(np.power(testData[:,-1] - treeMean, 2))
        if errorMerge < errorNoMerge:
            print('merging')
            return treeMean # 和并两个叶节点，返回合并后节点值
        else:
            return tree # 否则不合并，返回该子树
    else:
        return tree # 否则不合并，返回该子树

# 模型效果计较
# 线性叶子节点 预测计算函数 直接返回 树叶子节点 值
def regTreeEval(model, inDat):
    return float(model)


def modelTreeEval(model, inDat):
    n = np.shape(inDat)[1]
    X = np.mat(np.ones((1, n + 1)))  # 增加一列
    X[:, 1:n + 1] = inDat
    return float(X * model)  # 返回 值乘以 线性回归系数


# 树预测函数
def treeForeCast(tree, inData, modelEval=regTreeEval):
    if not isTree(tree):
        return modelEval(tree, inData)  # 返回 叶子节点 预测值
    if inData[tree['spInd']] > tree['spVal']:  # 左树
        if isTree(tree['left']):
            return treeForeCast(tree['left'], inData, modelEval)  # 还是树 则递归调用
        else:
            return modelEval(tree['left'], inData)  # 计算叶子节点的值 并返回
    else:
        if isTree(tree['right']):  # 右树
            return treeForeCast(tree['right'], inData, modelEval)
        else:
            return modelEval(tree['right'], inData)  # 计算叶子节点的值 并返回


# 得到预测值
def createForeCast(tree, testData, modelEval=regTreeEval):
    m = len(testData)
    yHat = np.mat(np.zeros((m, 1)))  # 预测标签
    for i in range(m):
        yHat[i, 0] = treeForeCast(tree, np.mat(testData[i]).tolist()[0], modelEval)
    return yHat


def getResult(yHat):
    resultList = []
    for i in range(len(yHat)):
        if yHat[i] > 0.5:
            resultList.append(1)
        else:
            resultList.append(0)
    return resultList



# 常量回归树和线性模型回归树的预测结果比较
def MRTvsSRT():
    TestMat = np.mat(loadDataSet('horseColicTest.txt'))
    TrainMat = np.mat(loadDataSet('horseColicTraining.txt'))
    # 普通回归树 预测结果
    # 得到普通回归树树
    # MyDatTest = loadDataSet('bikeSpeedVsIq_test.txt')
    # MyMatTest = np.mat(MyDatTest) # 是否对普通回归树进行后剪枝
    StaTree = createTree(TrainMat, ops=(1, 4))
    print(StaTree)
    # StaTree = prune(StaTree, MyMatTest)
    # 得到预测结果
    StaYHat = createForeCast(StaTree, TestMat[:, :-1], regTreeEval)  # 第一列为自变量!! 预测的是第二列的值
    resultList = getResult(StaYHat)
    print(resultList)
    err = 0
    m = len(resultList)
    for i in range(m):
        if resultList[i] != TestMat[i, -1]:
            err += 1
        # print(resultList[i], TestMat[i,3])
    print("正确率率是：%f" % (1 - err / m))
    # 预测结果和真实标签的相关系数
    # StaCorr = np.corrcoef(StaYHat, TestMat[:, -1], rowvar=0)[0, 1]  # NumPy 库函数
    # print(StaCorr)
    # 模型回归树 预测结果
    # 得到模型回归树
    # ModeTree = createTree(TrainMat, leafType=modelLeaf, errType=modelErr, ops=(1, 20))
    # # 得到预测结果
    # ModeYHat = createForeCast(ModeTree, TestMat[:, -1], modelTreeEval)
    # # 预测结果和真实标签的相关系数
    # ModeCorr = np.corrcoef(ModeYHat, TestMat[:, -1], rowvar=0)[0, 1]  # NumPy 库函数
    # print("普通回归树 预测结果的相关系数R2: %f" % (StaCorr))
    # print("模型回归树 预测结果的相关系数R2: %f" % (ModeCorr))
    # if ModeCorr > StaCorr:
    #     print("模型回归树效果优于普通回归树")
    # else:
    #     print("回归回归树效果优于模型普通树")


if __name__ == '__main__':
    MRTvsSRT()